package com.maben.config;

import org.slf4j.MDC;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.Executor;

/**
 * 自定义异步方法相关，添加MDC相关参数
 */
@Configuration
@EnableAsync
public class AsyncConfig {

    @Bean(name = "myAsyncExecutor")
    public Executor asyncExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(5);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(25);
        executor.setThreadNamePrefix("async-");
        
        // 设置任务装饰器，解决异步线程MDC丢失问题
        executor.setTaskDecorator(new MdcTaskDecorator());
        
        executor.initialize();
        return executor;
    }

    static class MdcTaskDecorator implements TaskDecorator {
        @Override
        public Runnable decorate(Runnable runnable) {
            // 捕获当前线程的MDC上下文
            Map<String, String> contextMap = MDC.getCopyOfContextMap();
            
            return () -> {
                // 在子线程中恢复MDC上下文
                if (contextMap != null) {
                    MDC.setContextMap(contextMap);
                }
                
                try {
                    // 执行被装饰的任务
                    runnable.run();
                } finally {
                    // 任务执行完成后清除MDC
                    MDC.clear();
                }
            };
        }
    }
}  